import requests
import time
import os

# Server URL
SERVER_URL = 'http://localhost:5000'

def start_inference(prompt, width=1024, height=1024, num_inference_steps=32, guidance_scale=5.0, num_images_per_prompt=1, seed=90):
    data = {
        'prompt': prompt,
        'width': width,
        'height': height,
        'num_inference_steps': num_inference_steps,
        'guidance_scale': guidance_scale,
        'num_images_per_prompt': num_images_per_prompt,
        'seed': seed,
        'task_id':int(time.time())
    }
    response = requests.post(f'{SERVER_URL}/start_inference', json=data)
    if response.status_code == 200:
        task_id = response.json().get('task_id')
        print(f'Task started with ID: {task_id}')
        return task_id
    else:
        print(f'Error starting task: {response.json()}')
        return None

def get_image(task_id):
    while True:
        response = requests.get(f'{SERVER_URL}/get_image/{task_id}/0')
        if response.status_code == 200:
            image_path = f'image_{task_id}.png'
            with open(image_path, 'wb') as f:
                f.write(response.content)
            print(f'Image saved as {image_path}')
            response2 = requests.get(f'{SERVER_URL}/get_image_info/{task_id}')
            print(response2)
            return image_path
        elif response.status_code == 202:
            print(f'Task {task_id} is still running. Checking again in 5 seconds...')
            time.sleep(5)
        else:
            print(f'Error fetching image: {response.json()}')
            return None

if __name__ == '__main__':
    prompt = "真人照片，岸边礁石有两个女生，左边女生在唱歌，右边女生弹吉他，两个人面对着镜头，镜头拉远，可以看到背后是大海，海那边是山，山在画面中间，还可以看到一点点夕阳"
    prompt ="原神可莉coser，真人照片"
    prompt= '布偶猫，在桌子上，布偶猫跟它的女主人玩耍'
    # Start inference task
    #task_id = start_inference(prompt)
    task_id = '1722099299'
    response2 = requests.get(f'{SERVER_URL}/get_image_info/{task_id}').json()
    print(response2)
    
    # if task_id:
    #     # Get the generated image
    #     get_image(task_id)
